package com.demo.function.shardingsphere.config;

import com.google.common.collect.Range;
import com.google.common.collect.Sets;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingValue;
import org.apache.shardingsphere.api.sharding.standard.RangeShardingValue;

import java.util.Collection;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.zip.CRC32;

/**
 * 复合分片策略，是精准分片和范围分片的组合
 * @author K
 */
@Slf4j
public class ComplexKeysShardingStrategy implements ComplexKeysShardingAlgorithm {

    @Override
    public Collection<String> doSharding(Collection availableTargetNames, ComplexKeysShardingValue shardingValue) {
        log.info("分片策略，数据库:{}，ShardingValue:{}", availableTargetNames, shardingValue);

        Map<String, Collection> columnNameAndShardingValuesMap = shardingValue.getColumnNameAndShardingValuesMap();
        Map<String, Range> columnNameAndRangeValuesMap = shardingValue.getColumnNameAndRangeValuesMap();

        // 先找精准分片字段，后找范围分片字段
        for (String column : ShardingConstant.COLUMNS) {
            if (columnNameAndShardingValuesMap.containsKey(column)) {
                Collection shardingValues = columnNameAndShardingValuesMap.get(column);
                return preciseDb(availableTargetNames, shardingValues);
            }

            if (columnNameAndRangeValuesMap.containsKey(column)) {
                Range range = columnNameAndRangeValuesMap.get(column);
                return rangDb(availableTargetNames, range);
            }
        }

        return availableTargetNames;
    }

    /**
     * 找出分片后的DB
     * @param availableTargetNames 表名集合
     * @param shardingValues 分片字段的值
     * @return db实例名集合
     */
    private Set<String> preciseDb(Collection availableTargetNames, Collection shardingValues) {
        Set<String> database = new HashSet<>();
        for (Object obj : shardingValues) {
            Long mod;
            if (obj instanceof Long) {
                Long tmp = (Long) obj;
                mod = tmp % ShardingConstant.DB_NUM;
            } else if (obj != null && obj instanceof String) {
                String tmp = (String) obj;
                CRC32 crc32 = new CRC32();
                crc32.update(tmp.getBytes());
                mod = crc32.getValue();
            } else {
                return Sets.newHashSet();
            }

            String tmpDbName = ShardingConstant.DB_NAME + mod;
            if (availableTargetNames.contains(tmpDbName)) {
                database.add(tmpDbName);
            }
        }
        log.info("分片后数据库：{}", database);
        return database;
    }

    /**
     * 根据值的范围找出相应DB
     * @param availableTargetNames 表名集合
     * @param range 范围值的对象
     * @return db实例名集合
     */
    private Set<String> rangDb(Collection availableTargetNames, Range range) {
        // 使用范围分片策略
        RangeShardingValue rangeShardingValue = new RangeShardingValue(null, null, range);
        RangeShardingStrategy rangeShardingStrategy = new RangeShardingStrategy();
        Collection collection = rangeShardingStrategy.doSharding(availableTargetNames, rangeShardingValue);
        return new HashSet<>(collection);
    }
}
